import json
import shutil
import sys

import torch
import time
import os
import cv2
from tqdm import *
from mAP.mAP import get_mAP_by_dir, get_mAP_by_program

# det = [['person', .88, 5, 67, 31, 48, '00001', (353, 500)],
#        ['person', .70, 119, 111, 40, 67, '00001', (353, 500)],
#        ['person', .80, 124, 9, 49, 67, '00001', (353, 500)]]
# gt = [['person', 25, 16, 38, 56, '00001', (353, 500)],
#       ['person', 129, 123, 41, 62, '00001', (353, 500)]]

resolutions = [240, 360, 480, 640, 720, 960, 1080, 1280]
currentPath = os.path.dirname(os.path.realpath(__file__))



def predict():
    # 使用模型预测并保存yolo格式的结果
    model = torch.hub.load('ultralytics/yolov5', 'custom', path='models/best.pt', device="cuda:0")
    for scene in range(1, 9):
        print("正在预测场景" + str(scene) + "的图像")

        image_path = os.path.join(currentPath, "datasets/AuAir/scene" + str(scene) + "/images/")
        image_name_list = os.listdir(image_path)
        for r in resolutions:
            predict_path = "datasets/AuAir/scene" + str(scene) + "/predict" + str(r) + "/"
            if not os.path.exists(predict_path):
                os.mkdir(predict_path)
            else:
                shutil.rmtree(predict_path)
                os.mkdir(predict_path)
            for i in tqdm(range(len(image_name_list))):
                results = model(os.path.join(image_path, image_name_list[i]), size=r)
                info = results.pandas().xyxy[0].to_json(orient="records")
                info = json.loads(info)
                file_name = image_name_list[i].split('.')[0] + '.txt'
                with open(os.path.join(predict_path, file_name), 'a', ) as file:
                    for item in info:
                        cls = item['name'].replace(' ', '-')
                        confidence = round(float(item["confidence"]), 4)
                        top = int(item['ymin'])
                        left = int(item['xmin'])
                        width = int(item['xmax']) - left
                        height = int(item['ymax']) - top
                        line = "{} {} {} {} {} {}\n".format(cls, confidence, left, top, width, height)
                        # line = "{} {} {} {} {}\n".format(cls, left, top, width, height)
                        file.write(line)


def predict2gt():
    for scene in tqdm(range(2, 9)):
        predict_dir = "datasets/AuAir/scene" + str(scene) + "/predict1280"
        gt_dir = "datasets/AuAir/scene" + str(scene) + "/anno1280"
        file_names = os.listdir(predict_dir)
        for file_name in file_names:
            with open(os.path.join(predict_dir, file_name), 'r', ) as file:
                info = file.readlines()
            with open(os.path.join(gt_dir, file_name), 'a', ) as file:
                for item in info:
                    item = item.replace("\n", '').split(' ')
                    line = "{} {} {} {} {}\n".format(item[0], item[2], item[3], item[4], item[5])
                    file.write(line)


def calc_mAP():
    # for scene in range(8, 9):
    #     print("场景" + str(scene) + "下的检测结果")
    #     gt_path = os.path.join(currentPath, "datasets/AuAir/scene" + str(scene) + "/anno1280")
    #     for resolution in resolutions:
    #         det_path = os.path.join(currentPath, "datasets/AuAir/scene" + str(scene) + "/predict" + str(resolution))
    #         result = get_mAP_by_dir(det_path, gt_path)
    #         print('分辨率为{}下的mAP值为: {} \n'.format(resolution, result))

    for scene in range(8, 9):
        print("场景" + str(scene) + "下的检测结果")
        gt_path = os.path.join(currentPath, "datasets/AuAir/scene" + str(scene) + "/anno1280")
        for resolution in resolutions:
            print("当前分辨率为{}px".format(resolution))
            for sample in range(1, 10):
                det_path = "datasets/AuAir/scene{}/predict{}_sample_{}".format(scene, resolution, sample / 10)
                det_path = os.path.join(currentPath, det_path)
                result = get_mAP_by_dir(det_path, gt_path)
                print('采样率为{}下的mAP值为: {} \n'.format(sample/10, result))


calc_mAP()
